#!/bin/bash

MODE=""
LIBTORCH_PATH=""
IPEX_VERSION=""

if [ $# == 3 ]; then
    MODE=$1
    LIBTORCH_PATH=$2
    IPEX_VERSION=$3
else
    echo "Usage: bash $0 [install|uninstall] <libtorch_path> <ipex_version>"
    exit
fi

if [[ ! ${MODE} == "install" ]] && [[ ! ${MODE} == "uninstall" ]]; then
    echo "Unrecognized execution mode: ${MODE}"
    exit
fi

if [ ! -f ${LIBTORCH_PATH}/build-version ]; then
    echo "${LIBTORCH_PATH} doesn't seem to be directory of libtorch."
    exit
fi

if [[ ${MODE} == "install" ]]; then
    TORCH_VERSION=$(cat ${LIBTORCH_PATH}/build-version)
    [[ ${TORCH_VERSION} =~ ([0-9]+\.[0-9]+).* ]]
    TORCH_VER=${BASH_REMATCH[1]}
    [[ ${IPEX_VERSION} =~ ([0-9]+\.[0-9]+).* ]]
    IPEX_VER=${BASH_REMATCH[1]}
    if [[ ${IPEX_VER} != ${TORCH_VER} ]]; then
        echo "Installation failed."
        echo "ERROR! Intel® Extension for PyTorch* needs to work with libtorch ${IPEX_VER}.*, but libtorch ${TORCH_VERSION} is found. Please switch to the matching version and run again."
        exit
    fi

    while IFS= read -r lib
    do
        if [ -f ${LIBTORCH_PATH}/lib/${lib} ]; then
            echo "Installation failed."
            echo "${lib} already exists in ${LIBTORCH_PATH}."
            exit
        fi
    done < <(ls -1 *.so)
    cp *.so ${LIBTORCH_PATH}/lib
    if [ $? -gt 0 ]; then
        echo "Installation failed."
        echo "Please check if ${LIBTORCH_PATH}/lib exists or if you have privileges to write in ${LIBTORCH_PATH}/lib."
        exit
    fi

    mkdir -p ${LIBTORCH_PATH}/share/cmake/intel_ext_pt_cpu
    if [ $? -gt 0 ]; then
        echo "Installation failed."
        echo "Please check if ${LIBTORCH_PATH}/share/cmake exists or if you have privileges to write in ${LIBTORCH_PATH}/share/cmake."
        exit
    fi
    cp intel_ext_pt_cpuConfig.cmake ${LIBTORCH_PATH}/share/cmake/intel_ext_pt_cpu

    echo $IPEX_VERSION > ${LIBTORCH_PATH}/build-version-libintel-ext-pt-cpu
    echo "Installation successed!"
elif [[ ${MODE} == "uninstall" ]]; then
    while IFS= read -r lib
    do
        if [ -f ${LIBTORCH_PATH}/lib/${lib} ]; then
            rm ${LIBTORCH_PATH}/lib/${lib}
        fi
    done < <(ls -1 *.so)
    if [ -d ${LIBTORCH_PATH}/share/cmake/intel_ext_pt_cpu ]; then
        rm -rf ${LIBTORCH_PATH}/share/cmake/intel_ext_pt_cpu
    fi
    if [ -f ${LIBTORCH_PATH}/build-version-libintel-ext-pt-cpu ]; then
        rm ${LIBTORCH_PATH}/build-version-libintel-ext-pt-cpu
    fi
    echo "Uninstallation successed!"
fi
